# Copyright 2023 Iguazio
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import json
from unittest.mock import MagicMock, patch

import pytest

import mlrun_pipelines
import mlrun_pipelines.common.helpers

import services.api.crud


def test_resolve_pipeline_project():
    cases = [
        {
            "expected_project": "project-from-deploy-p",
            "template": {
                "container": {
                    "command": [
                        "python",
                        "-m",
                        "mlrun",
                        "deploy",
                        "-p",
                        "project-from-deploy-p",
                    ]
                }
            },
        },
        {
            "expected_project": "project-from-deploy--project",
            "template": {
                "container": {
                    "command": [
                        "python",
                        "-m",
                        "mlrun",
                        "deploy",
                        "--project",
                        "project-from-deploy--project",
                    ]
                }
            },
        },
        {
            "expected_project": "project-from-deploy-f",
            "template": {
                "container": {
                    "command": [
                        "python",
                        "-m",
                        "mlrun",
                        "deploy",
                        "-f",
                        "db://project-from-deploy-f/tf2-serving@2db2ec7d89c0c8c9d1b9a86279d8440ebc230597",
                    ]
                }
            },
        },
        {
            "expected_project": "project-from-deploy--func-url",
            "template": {
                "container": {
                    "command": [
                        "python",
                        "-m",
                        "mlrun",
                        "deploy",
                        "--func-url",
                        "db://project-from-deploy--func-url/tf2-serving@2db2ec7d89c0c8c9d1b9a86279d8440ebc230597",
                    ]
                }
            },
        },
        {
            "expected_project": "project-from-deploy-precedence-p",
            "template": {
                "container": {
                    "command": [
                        "python",
                        "-m",
                        "mlrun",
                        "deploy",
                        "--func-url",
                        "db://project-from-deploy--func-url/tf2-serving@2db2ec7d89c0c8c9d1b9a86279d8440ebc230597",
                        "-p",
                        "project-from-deploy-precedence-p",
                    ]
                }
            },
        },
        {
            "expected_project": "project-from-run--project",
            "template": {
                "container": {
                    "command": [
                        "python",
                        "-m",
                        "mlrun",
                        "run",
                        "--project",
                        "project-from-run--project",
                    ]
                }
            },
        },
        {
            "expected_project": "project-from-run-f",
            "template": {
                "container": {
                    "command": [
                        "python",
                        "-m",
                        "mlrun",
                        "run",
                        "-f",
                        "db://project-from-run-f/tf2-serving@2db2ec7d89c0c8c9d1b9a86279d8440ebc230597",
                    ]
                }
            },
        },
        {
            "expected_project": "project-from-run--func-url",
            "template": {
                "container": {
                    "command": [
                        "python",
                        "-m",
                        "mlrun",
                        "run",
                        "--func-url",
                        "db://project-from-run--func-url/tf2-serving@2db2ec7d89c0c8c9d1b9a86279d8440ebc230597",
                    ]
                }
            },
        },
        {
            "expected_project": "project-from-run-r",
            "template": {
                "container": {
                    "command": [
                        "python",
                        "-m",
                        "mlrun",
                        "run",
                        "-r",
                        "{'kind': 'job', 'metadata': {'project': 'project-from-run-r'}}",
                    ]
                }
            },
        },
        {
            "expected_project": "project-from-run--runtime",
            "template": {
                "container": {
                    "command": [
                        "python",
                        "-m",
                        "mlrun",
                        "run",
                        "--runtime",
                        "{'kind': 'job', 'metadata': {'project': 'project-from-run--runtime'}}",
                    ]
                }
            },
        },
        {
            "expected_project": "project-from-run-precedence--project",
            "template": {
                "container": {
                    "command": [
                        "python",
                        "-m",
                        "mlrun",
                        "run",
                        "--func-url",
                        "db://project-from-deploy--func-url/tf2-serving@2db2ec7d89c0c8c9d1b9a86279d8440ebc230597",
                        "--project",
                        "project-from-run-precedence--project",
                    ]
                }
            },
        },
        {
            "expected_project": "project-from-build-r",
            "template": {
                "container": {
                    "command": [
                        "python",
                        "-m",
                        "mlrun",
                        "build",
                        "-r",
                        "{'kind': 'job', 'metadata': {'project': 'project-from-build-r'}}",
                    ]
                }
            },
        },
        {
            "expected_project": "project-from-build--runtime",
            "template": {
                "container": {
                    "command": [
                        "python",
                        "-m",
                        "mlrun",
                        "build",
                        "--runtime",
                        "{'kind': 'job', 'metadata': {'project': 'project-from-build--runtime'}}",
                    ]
                }
            },
        },
        {
            "expected_project": "project-from-build-precedence--runtime",
            "template": {
                "container": {
                    "command": [
                        "python",
                        "-m",
                        "mlrun",
                        "build",
                        "--runtime",
                        "{'kind': 'job', 'metadata': {'project': 'project-from-build--runtime'}}",
                        "--project",
                        "project-from-build-precedence--runtime",
                    ]
                }
            },
        },
        {
            "expected_project": "project-from-annotation",
            "template": {
                "metadata": {
                    "annotations": {
                        mlrun_pipelines.common.helpers.PROJECT_ANNOTATION: "project-from-annotation"
                    }
                }
            },
        },
    ]
    for case in cases:
        workflow_manifest = {"spec": {"templates": [case["template"]]}}
        pipeline = {
            "pipeline_spec": {"workflow_manifest": json.dumps(workflow_manifest)}
        }
        project = services.api.crud.Pipelines()._resolve_project_from_pipeline(
            mlrun_pipelines.models.PipelineRun(pipeline)
        )
        assert project == case["expected_project"]


@pytest.mark.parametrize(
    "project,expected_ids",
    [
        ("project-a", ["run1"]),
        ("*", ["run1", "run2"]),
    ],
)
def test_list_pipelines_project_filtering(project, expected_ids):
    pipelines = services.api.crud.pipelines.Pipelines()
    db_session = MagicMock()

    # Mock runs and KFP client
    run1 = MagicMock(id="run1", name="pipeline1", status="Succeeded")
    run2 = MagicMock(id="run2", name="pipeline2", status="Failed")
    all_runs = [run1, run2]
    mock_kfp_client = MagicMock()
    mock_kfp_client.list_runs.return_value = [(all_runs, None)]

    with (
        patch.object(
            services.api.crud.pipelines.Pipelines,
            "_initialize_kfp_client",
            return_value=mock_kfp_client,
        ),
        patch.object(
            services.api.crud.pipelines.Pipelines,
            "_resolve_project_from_pipeline",
            side_effect=lambda run: "project-a" if run.id == "run1" else "project-b",
        ),
        patch.object(
            services.api.crud.pipelines.Pipelines,
            "_format_runs_concurrently",
            side_effect=lambda kfp_client, runs, format_: [
                {"id": r.id, "name": r.name} for r in runs
            ],
        ),
    ):
        total_size, next_page_token, runs = pipelines.list_pipelines(
            db_session=db_session,
            project=project,
        )

    assert total_size == len(expected_ids)
    assert next_page_token is None
    assert [r["id"] for r in runs] == expected_ids
